# ai_detection/scene_classifier.py - Scene classification

import cv2
import numpy as np
import logging
from collections import Counter

logger = logging.getLogger(__name__)


class SceneClassifier:
    """
    Class for classifying football scenes
    """

    def __init__(self):
        # Initialize the YOLO detector
        from ai_detection.yolo_detector import YoloDetector
        self.yolo_detector = YoloDetector()

    def classify_scene(self, frames):
        """
        Classify a sequence of frames into one of the predefined categories

        Args:
            frames: List of frames (images)

        Returns:
            (scene_type, confidence) where scene_type is one of:
            - 'stadium' (wide shot of the field)
            - 'shot' (goal attempt)
            - 'celebration' (player/crowd celebration)
        """
        if not frames:
            return 'unknown', 0.0

        try:
            # Get scores for each scene type
            stadium_score = self.detect_stadium_shot(frames)
            shot_score = self.detect_goal_shot(frames)
            celebration_score = self.detect_celebration_shot(frames)

            # Determine the scene type with highest confidence
            scores = {
                'stadium': stadium_score,
                'shot': shot_score,
                'celebration': celebration_score
            }

            scene_type = max(scores, key=scores.get)
            confidence = scores[scene_type]

            return scene_type, confidence

        except Exception as e:
            logger.error(f"Error in scene classification: {str(e)}")
            return 'unknown', 0.0

    def detect_stadium_shot(self, frames):
        """
        Detect if frames show a stadium wide shot

        Returns:
            Confidence score (0-1)
        """
        if not frames:
            return 0.0

        try:
            # Stadium shots typically have:
            # 1. Green field covering a large portion of the frame
            # 2. Many small players (small bounding boxes)
            # 3. Horizontal field lines

            stadium_confidence = 0.0

            # Sample frames (use first, middle, last)
            sample_indices = [0, len(frames) // 2, -1]
            sample_frames = [frames[i] for i in sample_indices if 0 <= i < len(frames)]

            for frame in sample_frames:
                # 1. Check for green field
                green_ratio = self.calculate_green_ratio(frame)

                # 2. Check for player sizes
                players = self.yolo_detector.detect_players(frame)

                if players:
                    # Calculate average player size relative to frame
                    frame_area = frame.shape[0] * frame.shape[1]
                    player_sizes = []

                    for player in players:
                        bbox = player['bbox']
                        player_area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                        relative_size = player_area / frame_area
                        player_sizes.append(relative_size)

                    avg_player_size = sum(player_sizes) / len(player_sizes)

                    # Small players indicate wide shot
                    if avg_player_size < 0.01:  # 1% of frame or less
                        stadium_confidence += 0.3
                    elif avg_player_size < 0.02:  # 2% of frame or less
                        stadium_confidence += 0.2

                    # Many players visible indicates wide shot
                    if len(players) > 10:
                        stadium_confidence += 0.3
                    elif len(players) > 5:
                        stadium_confidence += 0.2

                # 3. Green field covering significant portion
                if green_ratio > 0.5:  # More than 50% green
                    stadium_confidence += 0.4
                elif green_ratio > 0.3:  # More than 30% green
                    stadium_confidence += 0.2

            # Calculate final confidence
            stadium_confidence = min(1.0, stadium_confidence / len(sample_frames))

            return stadium_confidence

        except Exception as e:
            logger.error(f"Error in stadium shot detection: {str(e)}")
            return 0.0

    def detect_goal_shot(self, frames):
        """
        Detect if frames show a goal attempt

        Returns:
            Confidence score (0-1)
        """
        if not frames:
            return 0.0

        try:
            # Let YOLO detector determine if this is a goal
            goal_confidence = self.yolo_detector.detect_goal(frames)

            # Additional analysis for shot detection
            shot_indicators = 0

            # Sample frames
            sample_indices = [0, len(frames) // 2, -1]
            sample_frames = [frames[i] for i in sample_indices if 0 <= i < len(frames)]

            for frame in sample_frames:
                # Check for ball
                ball = self.yolo_detector.detect_ball(frame)
                if ball:
                    shot_indicators += 1

                # Check for goalie/goalpost in the frame
                objects = self.yolo_detector.detect_objects(frame)
                for obj in objects:
                    if obj['class_name'] == 'sports':
                        shot_indicators += 0.5

            # Normalize shot indicators
            shot_indicator_score = min(1.0, shot_indicators / (len(sample_frames) * 2))

            # Combine scores
            combined_score = max(goal_confidence, shot_indicator_score)

            return combined_score

        except Exception as e:
            logger.error(f"Error in goal shot detection: {str(e)}")
            return 0.0

    def detect_celebration_shot(self, frames):
        """
        Detect if frames show a celebration

        Returns:
            Confidence score (0-1)
        """
        if not frames:
            return 0.0

        try:
            # Let YOLO detector calculate celebration probability
            celebration_confidence = self.yolo_detector.detect_celebration(frames)

            # Check visual characteristics of celebration
            # - Camera may shake or move quickly
            # - Close-up of players
            # - Crowd animation

            motion_scores = []
            player_size_scores = []

            # Calculate motion between consecutive frames
            for i in range(len(frames) - 1):
                if i % 5 == 0:  # Sample every 5 frames to save processing
                    motion = self.calculate_motion(frames[i], frames[i + 1])
                    motion_scores.append(motion)

            # Check player sizes (close-ups)
            for frame in frames:
                players = self.yolo_detector.detect_players(frame)

                if players:
                    # Calculate average player size relative to frame
                    frame_area = frame.shape[0] * frame.shape[1]
                    player_sizes = []

                    for player in players:
                        bbox = player['bbox']
                        player_area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                        relative_size = player_area / frame_area
                        player_sizes.append(relative_size)

                    avg_player_size = sum(player_sizes) / len(player_sizes)

                    # Larger players indicate close-up (celebration)
                    if avg_player_size > 0.1:  # 10% of frame or more
                        player_size_scores.append(1.0)
                    elif avg_player_size > 0.05:  # 5% of frame or more
                        player_size_scores.append(0.7)
                    else:
                        player_size_scores.append(0.0)

            # Combine scores
            motion_score = np.mean(motion_scores) if motion_scores else 0
            player_size_score = np.mean(player_size_scores) if player_size_scores else 0

            combined_score = (celebration_confidence * 0.6 +
                              motion_score * 0.2 +
                              player_size_score * 0.2)

            return combined_score

        except Exception as e:
            logger.error(f"Error in celebration shot detection: {str(e)}")
            return 0.0

    def calculate_green_ratio(self, frame):
        """Calculate ratio of green pixels in the frame (for field detection)"""
        try:
            # Convert to HSV for better color detection
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

            # Define range of green color in HSV
            lower_green = np.array([35, 50, 50])
            upper_green = np.array([85, 255, 255])

            # Create a mask for green color
            mask = cv2.inRange(hsv, lower_green, upper_green)

            # Count green pixels
            green_pixels = cv2.countNonZero(mask)
            total_pixels = frame.shape[0] * frame.shape[1]

            return green_pixels / total_pixels

        except Exception as e:
            logger.error(f"Error in green ratio calculation: {str(e)}")
            return 0.0

    def calculate_motion(self, frame1, frame2):
        """Calculate motion between two consecutive frames"""
        try:
            # Convert to grayscale
            gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
            gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

            # Calculate optical flow
            flow = cv2.calcOpticalFlowFarneback(
                gray1, gray2, None, 0.5, 3, 15, 3, 5, 1.2, 0)

            # Calculate magnitude of flow
            mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])

            # Normalize motion score
            motion_score = np.mean(mag) / 10.0  # Normalize to 0-1 range
            motion_score = min(1.0, motion_score)

            return motion_score

        except Exception as e:
            logger.error(f"Error in motion calculation: {str(e)}")
            return 0.0